// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux::test;

namespace kumo::pollux::functions {
namespace {

class MapTopNValuesTest : public test::FunctionBaseTest {};

TEST_F(MapTopNValuesTest, emptyMap) {
  RowVectorPtr input = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{}",
      }),
  });

  assertEqualVectors(
      evaluate("map_top_n_values(c0, 3)", input),
      make_array_vector_from_json<int64_t>({
          "[]",
      }));
}

TEST_F(MapTopNValuesTest, basic) {
  auto data = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{1:3, 2:5, 3:1, 4:4, 5:2}",
          "{1:3, 2:5, 3:null, 4:4, 5:2}",
          "{1:null, 2:null, 3:1, 4:4, 5:null}",
          "{1:10, 2:7, 3:11, 5:4}",
          "{1:10, 2:10, 3:10, 4:10, 5:10}",
          "{1:10, 2:7, 3:0}",
          "{1:null, 2:10}",
          "{}",
          "{1:null, 2:null, 3:null}",
      }),
  });

  auto result = evaluate("map_top_n_values(c0, 3)", data);

  auto expected = make_array_vector_from_json<int64_t>({
      "[5, 4, 3]",
      "[5, 4, 3]",
      "[4, 1, null]",
      "[11, 10, 7]",
      "[10, 10, 10]",
      "[10, 7, 0]",
      "[10, null]",
      "[]",
      "[null, null, null]",
  });

  assertEqualVectors(expected, result);

  // n = 0. Expect empty maps.
  result = evaluate("map_top_n_values(c0, 0)", data);

  expected = make_array_vector_from_json<int64_t>({
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
      "[]",
  });

  assertEqualVectors(expected, result);

  // n is negative. Expect an error.
  POLLUX_ASSERT_THROW(
      evaluate("map_top_n_values(c0, -1)", data),
      "n must be greater than or equal to 0");
}

TEST_F(MapTopNValuesTest, complexKeys) {
  RowVectorPtr input =
      make_row_vector({make_map_vector_from_json<std::string, std::string>(
          {R"( {"x":"ab", "y":"y"} )",
           R"( {"x":"dw", "x2":"-2"} )",
           R"( {"ac":"ac", "cc":"cc", "dd": "dd"} )"})});

  assertEqualVectors(
      evaluate("map_top_n_values(c0, 1)", input),
      make_array_vector_from_json<std::string>({
          "[\"y\"]",
          "[\"dw\"]",
          "[\"dd\"]",
      }));
}

TEST_F(MapTopNValuesTest, tryFuncWithInputHasNullInArray) {
  auto array_vector = make_nullable_flat_vector<std::int32_t>({1, 2, 3});
  auto flat_vector = make_nullable_array_vector<bool>(
      {{true, std::nullopt}, {true, std::nullopt}, {false, std::nullopt}});
  auto mapvector = make_map_vector(
      /*offsets=*/{0},
      /*keyVector=*/array_vector,
      /*value_vector=*/flat_vector);
  auto rst = evaluate(
      "try(map_top_n_values(c0, 6455219767830808341))",
      make_row_vector({mapvector}));
  assert(rst->size() == 1);
  assert(rst->is_null_at(0));
}

TEST_F(MapTopNValuesTest, floatingPointValues) {
  // Test with floating point values
  auto data = make_row_vector({
      make_map_vector_from_json<int32_t, double>({
          "{1:3.5, 2:5.1, 3:1.2, 4:4.9, 5:2.0}",
          "{1:3.5, 2:5.1, 3:null, 4:4.9, 5:2.0}",
      }),
  });

  auto result = evaluate("map_top_n_values(c0, 3)", data);

  auto expected = make_array_vector_from_json<double>({
      "[5.1, 4.9, 3.5]",
      "[5.1, 4.9, 3.5]",
  });

  assertEqualVectors(expected, result);
}

TEST_F(MapTopNValuesTest, edgeCaseValues) {
  // Test with minimum/maximum values
  auto data = make_row_vector({
      make_map_vector_from_json<int32_t, int64_t>({
          "{1:" + std::to_string(INT64_MIN) +
              ", 2:" + std::to_string(INT64_MAX) + "}",
          "{1:0, 2:" + std::to_string(INT64_MAX) + "}",
      }),
  });

  auto result = evaluate("map_top_n_values(c0, 2)", data);

  auto expected = make_array_vector_from_json<int64_t>({
      "[" + std::to_string(INT64_MAX) + ", " + std::to_string(INT64_MIN) + "]",
      "[" + std::to_string(INT64_MAX) + ", 0]",
  });

  assertEqualVectors(expected, result);
}

TEST_F(MapTopNValuesTest, stringValuesWithSpecialChars) {
  // Test string values with special characters
  auto data = make_row_vector({
      make_map_vector_from_json<int32_t, std::string>({
          R"({1:"a,b", 2:"a\nb", 3:"a\tb"})",
          R"({1:"", 2:" ", 3:"  "})",
      }),
  });

  auto result = evaluate("map_top_n_values(c0, 2)", data);

  auto expected = make_array_vector_from_json<std::string>({
      R"(["a,b", "a\nb"])",
      R"(["  ", " "])",
  });

  assertEqualVectors(expected, result);
}
} // namespace
} // namespace kumo::pollux::functions
